import torch

from .builder import EncoderDecoder
from .head import SegFormerHead
from .backbone import MiT

segformer_settings = {
    'B0': 256,  # head_dim
    'B1': 256,
    'B2': 768,
    'B3': 768,
    'B4': 768,
    'B5': 768
}

mit_settings = {
    'B0': [[32, 64, 160, 256], [2, 2, 2, 2]],  # [embed_dims, depths]
    'B1': [[64, 128, 320, 512], [2, 2, 2, 2]],
    'B2': [[64, 128, 320, 512], [3, 4, 6, 3]],
    'B3': [[64, 128, 320, 512], [3, 4, 18, 3]],
    'B4': [[64, 128, 320, 512], [3, 8, 27, 3]],
    'B5': [[64, 128, 320, 512], [3, 6, 40, 3]]
}


class SegFormer(EncoderDecoder):
    def __init__(self, variant: str = 'B0', num_classes: int = 19):
        in_channels = mit_settings[variant][0]
        backbone = MiT(variant)
        decode_head = SegFormerHead(dims=in_channels, embed_dim=segformer_settings[variant],
                                    num_classes=num_classes)

        super(SegFormer, self).__init__(backbone=backbone,
                                        decode_head=decode_head,
                                        in_channels=in_channels)


if __name__ == '__main__':
    x = torch.randn(2, 3, 224, 224)
    model = SegFormer("B0", 19)

    y = model(x)
    print(y.shape)
